import "@moonbeam-network/api-augment";
import {
  MoonwallContext,
  beforeAll,
  describeSuite,
  expect,
  type ChopsticksContext,
} from "@moonwall/cli";
import type { ApiPromise } from "@polkadot/api";
import type { HexString } from "@polkadot/util/types";
import type { u32 } from "@polkadot/types";
import { hexToU8a, u8aConcat, u8aToHex } from "@polkadot/util";
import { blake2AsHex, xxhashAsU8a } from "@polkadot/util-crypto";
import { env } from "node:process";

// Index of the `MigrateDelegationScheduledRequestsToDoubleMap` migration in
// the `MultiBlockMigrations` tuple defined in `runtime/common/src/migrations.rs`.
// This is shared by Moonbase, Moonriver and Moonbeam.
const STAKING_MIGRATION_INDEX = 2;

const hash = (prefix: HexString, suffix: Uint8Array) => {
  return u8aToHex(u8aConcat(hexToU8a(prefix), xxhashAsU8a(suffix, 64), suffix));
};

const upgradeRestrictionSignal = (paraId: u32) => {
  const prefix = "0xcd710b30bd2eab0352ddcc26417aa194f27bbb460270642b5bcaf032ea04d56a";

  return hash(prefix, paraId.toU8a());
};

const upgradeRuntime = async (context: ChopsticksContext) => {
  const path = (await MoonwallContext.getContext()).rtUpgradePath;
  if (!path) {
    throw new Error("Runtime wasm path (rtUpgradePath) is not configured");
  }

  const rtWasm = await import("node:fs").then(({ readFileSync }) => readFileSync(path));
  const rtHex = `0x${rtWasm.toString("hex")}`;
  const rtHash = blake2AsHex(rtHex);
  const api = context.polkadotJs();
  const signer = context.keyring.alice;

  await context.setStorage({
    module: "system",
    method: "authorizedUpgrade",
    methodParams: `${rtHash}01`, // 01 is for the RT ver check = true
  });
  await context.createBlock();

  await api.tx.system.applyAuthorizedUpgrade(rtHex).signAndSend(signer);

  const paraId: u32 = (await api.query.parachainInfo.parachainId()) as unknown as u32;

  await api.rpc("dev_newBlock", {
    count: 3,
    relayChainStateOverrides: [[upgradeRestrictionSignal(paraId), null]],
  });
};

describeSuite({
  id: "C03",
  title: "Chopsticks Staking Migration - DelegationScheduledRequests",
  foundationMethods: "chopsticks",
  testCases: ({ it, context, log }) => {
    let api: ApiPromise;
    let specName: string;

    beforeAll(async () => {
      api = context.polkadotJs();
      await api.isReady;

      specName = (api.consts.system.version as any).specName.toString();
      log(`Connected to chain specName=${specName}`);
    });

    it({
      id: "T1",
      timeout: 900_000,
      title:
        "DelegationScheduledRequests is migrated from single map to double map and counters initialized",
      test: async () => {
        const isMoonbeam = specName === "moonbeam";
        const psQueryBefore: any = api.query.parachainStaking;

        // 1. Capture the pre-upgrade DelegationScheduledRequests layout (single map).
        const oldEntries = await psQueryBefore.delegationScheduledRequests.entries();
        let totalOldRequests = 0;
        const expectedPerCollator: Record<string, number> = {};
        for (const [storageKey, boundedVec] of oldEntries as any) {
          const requestsJson = (boundedVec as any).toJSON() as any[];
          totalOldRequests += requestsJson.length;

          // Derive per collator how many distinct
          // delegators have at least one scheduled request before the
          // migration. This is the expected target for
          // `DelegationScheduledRequestsPerCollator` after migration.
          const collator = (storageKey as any).args?.[0]?.toString?.() ?? "";
          if (collator) {
            const uniqueDelegators = new Set<string>();
            for (const req of requestsJson as any[]) {
              const delegator = (req as any)?.delegator;
              if (delegator != null) {
                uniqueDelegators.add(String(delegator));
              }
            }
            expectedPerCollator[collator] = uniqueDelegators.size;
          }
        }

        log(`Pre-upgrade DelegationScheduledRequests entries (requests): ${totalOldRequests}`);

        const rtBefore = (api.consts.system.version as any).specVersion.toNumber();
        log(`Spec version before upgrade: ${rtBefore}`);

        // 2. Perform the runtime upgrade which includes the new multi-block migration.
        await upgradeRuntime(context as unknown as ChopsticksContext);

        const rtAfter = (api.consts.system.version as any).specVersion.toNumber();
        log(`Spec version after upgrade: ${rtAfter}`);
        expect(rtAfter).to.be.greaterThan(rtBefore);

        // Wait briefly so the API can refresh metadata after the upgrade.
        await new Promise((resolve) => setTimeout(resolve, 1_000));
        await api.isReady;

        // Helper to read the current staking migration state in a single block.
        const readState = async () => {
          const psQueryAfter: any = api.query.parachainStaking;

          // Always read the per-collator counters. On Moonbeam we avoid
          // scanning the entire new double-map and rely solely on these
          // counters, while on smaller networks we additionally read the
          // double-map to enforce a stronger invariant.
          const perCollatorCounterQuery: any = psQueryAfter.delegationScheduledRequestsPerCollator;
          const counters: Record<string, number> = {};
          let totalDelegatorQueues = 0;
          if (perCollatorCounterQuery && perCollatorCounterQuery.entries) {
            const counterEntries = await perCollatorCounterQuery.entries();
            for (const [storageKey, count] of counterEntries as any) {
              const collator = (storageKey as any).args?.[0]?.toString?.() ?? "";
              if (collator) {
                const value = (count as any).toNumber();
                counters[collator] = value;
                totalDelegatorQueues += value;
              }
            }
          }

          if (isMoonbeam) {
            return {
              counters,
              totalDelegatorQueues,
            };
          }

          const entries = await psQueryAfter.delegationScheduledRequests.entries();
          let totalRequests = 0;
          for (const [, boundedVec] of entries as any) {
            const requestsJson = (boundedVec as any).toJSON() as any[];
            totalRequests += requestsJson.length;
          }

          return {
            counters,
            totalRequests,
            queueCount: entries.length,
            totalDelegatorQueues,
          };
        };

        // 3. Progress blocks while the multi-block migrations are running and
        //    assert consistency. In full mode we check invariants on every
        //    block while the *staking* migration is active and once after it
        //    completes. In light mode we only read staking storage and assert
        //    invariants once the staking migration has completed.
        //
        // NOTE: The multi-block migration framework lives in the
        // `multiBlockMigrations` pallet. Its `cursor` storage item returns
        // `Option<PalletMigrationsMigrationCursor>`, where:
        // - `None`   => no migration is currently running
        // - `Some`   => an active migration `{ active: { index, innerCursor, startedAt } }`
        //              or a stuck migration.
        const migrationsQuery: any = (api.query as any).multiBlockMigrations;
        let blocksAfterUpgrade = 0;
        let sawStakingMigration = false;

        // Always check at least one block after the upgrade, then keep going
        // until `pallet-migrations` cursor becomes `None` or we hit a hard cap.
        for (let i = 0; i < 128; i++) {
          await context.createBlock();
          blocksAfterUpgrade += 1;

          const cursor = migrationsQuery?.cursor ? await migrationsQuery.cursor() : null;
          const cursorStr = cursor?.toString?.() ?? "n/a";

          // Decode the active migration index from the cursor, if any.
          const cursorJson = cursor?.toJSON?.() as any;
          const activeIndex: number | null =
            cursorJson && cursorJson.active && typeof cursorJson.active.index === "number"
              ? (cursorJson.active.index as number)
              : null;

          // Track when the staking migration has been seen at least once.
          if (activeIndex === STAKING_MIGRATION_INDEX) {
            sawStakingMigration = true;
          }

          const isCursorNone = !!cursor && (cursor as any).isNone;

          // Only observe cursor progress until we know the staking migration
          // has finished, then take a single state snapshot and assert
          // invariants. This keeps the test lightweight even on large chains.
          if (!isCursorNone) {
            log(`Block +${blocksAfterUpgrade}: cursor=${cursorStr}`);
            continue;
          }

          // cursor is None here: all multi-block migrations have finished.
          // If we never saw the staking migration become active something
          // went wrong with the configuration.
          if (!sawStakingMigration) {
            throw new Error(
              "Staking migration did not appear in multiBlockMigrations cursor before completion"
            );
          }

          const state = (await readState()) as any;
          const counters = state.counters as Record<string, number>;

          // On all networks, assert that for every collator that had at least
          // one scheduled request before the migration, the final per-collator
          // counter matches the number of delegators with at least one
          // scheduled request.
          for (const [collator, expectedCount] of Object.entries(expectedPerCollator)) {
            const actual = counters[collator] ?? 0;
            expect(actual).to.equal(
              expectedCount,
              `DelegationScheduledRequestsPerCollator[${collator}] must equal number of delegators with at least one scheduled request before migration`
            );
          }

          if (!isMoonbeam) {
            const { totalRequests, queueCount, totalDelegatorQueues } = state;

            log(
              `Block +${blocksAfterUpgrade}: totalRequests=${totalRequests}, queues=${queueCount}, sumCounters=${totalDelegatorQueues}, cursor=${cursorStr}`
            );

            expect(totalRequests).to.equal(
              totalOldRequests,
              "Total number of scheduled delegation requests must be preserved during migration"
            );

            if (queueCount > 0) {
              expect(totalDelegatorQueues).to.equal(
                queueCount,
                "Sum of DelegationScheduledRequestsPerCollator values should equal number of (collator, delegator) queues after migration completes"
              );
            }
          }
          break;
        }
      },
    });
  },
});
